#include <Adafruit_GFX.h>
#include <Adafruit_ST7735.h>
#include <SPI.h>

#include <TensorFlowLite.h>
#include <tensorflow/lite/micro/all_ops_resolver.h>
#include <tensorflow/lite/micro/micro_error_reporter.h>
#include <tensorflow/lite/micro/micro_interpreter.h>
#include <tensorflow/lite/schema/schema_generated.h>
#include <tensorflow/lite/version.h>

#include "Arduino_BMI270_BMM150.h"
#include "model3_60_60.h"

const float accelerationThreshold = 2.5;
const int numSamples = 119;
int samplesRead = numSamples;

#define TFT_CS		10
#define TFT_RST		8
#define TFT_DC		9

Adafruit_ST7735 tft = Adafruit_ST7735(TFT_CS, TFT_DC, TFT_RST);

float p = 3.1415926;
tflite::MicroErrorReporter tflErrorReporter;
tflite::AllOpsResolver tflOpsResolver;
TfLiteTensor* tflInputTensor = nullptr;
TfLiteTensor* tflOutputTensor = nullptr;
const tflite::Model* tflModel = nullptr;
tflite::MicroInterpreter* tflInterpreter = nullptr;

constexpr int tensorArenaSize = 8 * 1024;
byte tensorArena[tensorArenaSize] __attribute__((aligned(16)));

const char* GESTURES[] = {
	"punch",
	"flex",
	"roll"
};

#define NUM_GESTURES (sizeof(GESTURES) / sizeof(GESTURES[0]))

void setup(void) {
	Serial.begin(9600);	
	tft.initR(INITR_BLACKTAB);
	Serial.println("TFT Initialized");

	uint16_t time = millis();
	tft.fillScreen(ST77XX_BLACK);
	time = millis() - time;
	Serial.println(time, DEC);
	testdrawtext("HELLO WORLD, MADE BY GANYIHUI", ST77XX_WHITE);

	if (!IMU.begin()) {
		Serial.println("Failed to initialize IMU!");
		while (1);
	}

	tflModel = tflite::GetModel(model);
	if (tflModel->version() != TFLITE_SCHEMA_VERSION) {
		Serial.println("Model schema mismatch!");
		while (1);
	}

	tflInterpreter = new tflite::MicroInterpreter(tflModel, tflOpsResolver, tensorArena, tensorArenaSize, &tflErrorReporter);
	tflInterpreter->AllocateTensors();

	tflInputTensor = tflInterpreter->input(0);
	tflOutputTensor = tflInterpreter->output(0);

	Serial.println("ML done");
}

void loop() {
	float aX, aY, aZ, gX, gY, gZ;
	while (samplesRead == numSamples) {
		if (IMU.accelerationAvailable()) {
			IMU.readAcceleration(aX, aY, aZ);
			if ((fabs(aX) + fabs(aY) + fabs(aZ)) >= accelerationThreshold) {
				samplesRead = 0;
				break;
			}
		}
	}
	while (samplesRead < numSamples) {
		if (IMU.accelerationAvailable() && IMU.gyroscopeAvailable()) {
			IMU.readAcceleration(aX, aY, aZ);
			IMU.readGyroscope(gX, gY, gZ);

			tflInputTensor->data.f[samplesRead * 6 + 0] = (aX + 4.0) / 8.0;
			tflInputTensor->data.f[samplesRead * 6 + 1] = (aY + 4.0) / 8.0;
			tflInputTensor->data.f[samplesRead * 6 + 2] = (aZ + 4.0) / 8.0;
			tflInputTensor->data.f[samplesRead * 6 + 3] = (gX + 2000.0) / 4000.0;
			tflInputTensor->data.f[samplesRead * 6 + 4] = (gY + 2000.0) / 4000.0;
			tflInputTensor->data.f[samplesRead * 6 + 5] = (gZ + 2000.0) / 4000.0;

			samplesRead++;

			if (samplesRead == numSamples) {
				// Run inferencing
				TfLiteStatus invokeStatus = tflInterpreter->Invoke();
				if (invokeStatus != kTfLiteOk) {
					Serial.println("Invoke failed!");
					while (1);
					return;
				}

				// Loop through the output tensor values from the model3
				float maxValue = 0.0;
				int maxI = 0;
				for (int i = 0; i < NUM_GESTURES; i++) {
					Serial.print(GESTURES[i]);
					Serial.println(": ");
					Serial.println(tflOutputTensor->data.f[i], 6);
					Serial.println();
					if (tflOutputTensor->data.f[i] > maxValue) {
						maxValue = tflOutputTensor->data.f[i];
						maxI = i;
					}
				}
				tft.fillScreen(ST77XX_BLACK);
				testdrawtext(const_cast<char*>(GESTURES[maxI]), ST77XX_WHITE);
			}
		}
	}
}

void testdrawtext(char *text, uint16_t color) {
	tft.setCursor(0, 0);
	tft.setTextColor(color);
	tft.setTextWrap(true);
	tft.print(text);
}